package cn.edu.njust.limiter;

import java.util.concurrent.*;

public class TokenBucketRateLimiter implements RateLimiter {
    private Semaphore semaphore;

    // 加令牌的时间间隔
    private int fuelInterval;

    // 令牌桶大小
    private int bucketSize;

    private ScheduledExecutorService scheduler;

    public TokenBucketRateLimiter(int bucketSize, int fuelInterval) {
        this.bucketSize = bucketSize;
        this.fuelInterval = fuelInterval;
        this.semaphore = new Semaphore(bucketSize);
        this.scheduler = Executors.newSingleThreadScheduledExecutor(new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                return new Thread(r, "CounterRateLimiterThread");
            }
        });
        // 按恒定速率往令牌桶中加令牌
        this.scheduler.scheduleAtFixedRate(() -> {
            if (semaphore.availablePermits() < bucketSize)
                semaphore.release(1);
        }, 0, fuelInterval, TimeUnit.MILLISECONDS);
    }

    @Override
    public void aquire() throws InterruptedException {
        semaphore.acquire();
    }
}
